﻿using System;
using Org.BouncyCastle.Crypto;
using Org.BouncyCastle.Crypto.Engines;
using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Security;
using System.IO;

namespace SafeFile.Ciphers
{
    internal class AES
    {
        private const short KEY_LENGTH = 256;

        public class KeyInfo
        {
            public byte[] Key;
            public byte[] IV;

            public byte[] GetBytes()
            {
                using (var stream = new MemoryStream())
                {
                    stream.Write(Key, 0, Key.Length);
                    stream.Write(IV, 0, IV.Length);
                    return stream.ToArray();
                }
            }
            public void SetBytes(byte[] data)
            {
                var keyLen = KEY_LENGTH / 8;
                Key = new byte[keyLen];
                Array.Copy(data, 0, Key, 0, keyLen);

                IV = new byte[16];
                Array.Copy(data, keyLen, IV, 0, 16);
            }
        }

        public static KeyInfo MakeKeys()
        {
            var keyInfo = new KeyInfo();

            // 生成256位的Aes key
            CipherKeyGenerator generator = GeneratorUtilities.GetKeyGenerator("AES");
            generator.Init(new KeyGenerationParameters(new SecureRandom(), KEY_LENGTH));
            keyInfo.Key = generator.GenerateKey();
            keyInfo.IV = new SecureRandom().GenerateSeed(16);

            return keyInfo;
        }

        public static void Encrypt(FileStream inputStream, FileStream outputStream, KeyInfo keyInfo, Action<int> progressReporter)
        {
            var engine = new AesEngine();
            var keyParam = new KeyParameter(keyInfo.Key);
            //ParametersWithIV keyParamWithIV = new ParametersWithIV(keyParam, keyInfo.IV);

            //engine.Init(true, keyParamWithIV);
            engine.Init(true, keyParam);

            var blockSize = engine.GetBlockSize();
            var handled = 0;

            while (true)
            {
                byte[] buffer = new byte[blockSize];

                var read = inputStream.Read(buffer, 0, blockSize);
                if (read == 0)
                {
                    break;
                }
                handled += read;
                progressReporter(handled);
                byte[] output = new byte[blockSize];
                var count = engine.ProcessBlock(buffer, 0, output, 0);
                outputStream.Write(output, 0, count);
            }
        }

        public static void Decrypt(FileStream inputStream, FileStream outputStream, KeyInfo keyInfo, int length, Action<int> progressReporter)
        {
            var engine = new AesEngine();
            var keyParam = new KeyParameter(keyInfo.Key);
            //ParametersWithIV keyParamWithIV = new ParametersWithIV(keyParam, keyInfo.IV);

            //engine.Init(true, keyParamWithIV);
            engine.Init(false, keyParam);

            var blockSize = engine.GetBlockSize();

            var handleLength = 0;

            while (true)
            {
                byte[] buffer = new byte[blockSize];

                var read = inputStream.Read(buffer, 0, blockSize);
                if (read == 0)
                {
                    break;
                }
                byte[] output = new byte[blockSize];
                var count = engine.ProcessBlock(buffer, 0, output, 0);


                if (handleLength + count <= length)
                {
                    handleLength += count;
                    progressReporter(handleLength);
                    outputStream.Write(output, 0, count);
                    continue;
                }

                // 处理的数据比实际数据更长
                if (inputStream.Position == inputStream.Length)
                {
                    // 没有数据了
                    count = length - handleLength;
                    outputStream.Write(output, 0, count);
                    break;
                }
                // 文件大小不正确
                throw new Exception("Invalid file.");
            }
            progressReporter(handleLength);
        }
    }
}
